# -*- coding: utf-8 -*-
import functools

import allure


def assertDecorator(f):
    @functools.wraps(f)
    def wrapper(self, *args, **kwargs):
        from PageClass import BasePage
        driver = args[-1]
        data = args[-2]
        try:
            f(self, *args, **kwargs)
        except AssertionError as e:
            imgPath = BasePage(driver).screenshot("error", data.caseName)
            allure.attach.file(imgPath, data.caseName, allure.attachment_type.PNG)
            raise e
        else:
            imgPath = BasePage(driver).screenshot("success", data.caseName)
            allure.attach.file(imgPath, data.caseName, allure.attachment_type.PNG)

    return wrapper


class ZTCCloudAssertion(object):

    @assertDecorator
    def assertIn(self, args1, args2, *args):
        try:
            assert args1 in args2
        except AssertionError:
            raise AssertionError(f'{args1} not in {args2}')

    @assertDecorator
    def assertEqual(self, args1, args2, *args):
        try:
            assert args1 == args2
        except AssertionError:
            raise AssertionError(f'{args1} not equal {args2}')


ZTCCloudAssertion = ZTCCloudAssertion
